#include "Precompiled.hpp"
#include "output.hpp"
#include <string>
using namespace std;


string get_cstring(ifstream& ifs, int offset)
{
    int pos = ifs.tellg();
    ifs.seekg(offset, ios_base::beg);
    string ret;
    char one_char;
    ifs.read(&one_char, 1);
    while (one_char != 0){
        ret += one_char;
        ifs.read(&one_char, 1);
    }
    ifs.seekg(pos, ios_base::beg);
    return ret;
}

template <class DataStructure>
vector<DataStructure> get_vector(ifstream& ifs, int offset)
{
    const int SIZE = sizeof(DataStructure);
    
    DataStructure empty_end;
    memset( &empty_end, 0, SIZE);
    int pos = ifs.tellg();
    ifs.seekg(offset, ios_base::beg);

    vector<DataStructure> ret;
    DataStructure one_data;
    memset(&one_data, 0, SIZE);
    ifs.read(reinterpret_cast<char*>(&one_data), SIZE);
    //Need consider EOF???
    while( memcmp (&one_data, &empty_end, SIZE) != 0){
        ret.push_back(one_data);
        memset(&one_data, 0, SIZE);
        ifs.read(reinterpret_cast<char*>(&one_data), SIZE);
    }
    ifs.seekg(pos, ios_base::beg);
    return ret;
}

template <class DataStructure>
DataStructure get_data(ifstream& ifs, int offset)
{
    const int SIZE = sizeof(DataStructure);
    
    int pos = ifs.tellg();
    ifs.seekg(offset, ios_base::beg);

    DataStructure one_data;
    memset(&one_data, 0, SIZE);
    ifs.read(reinterpret_cast<char*>(&one_data), SIZE);

    ifs.seekg(pos, ios_base::beg);
    return one_data;
}


int main()
{
    cout << "*******************************************************" << endl;
    cout << "****            dump PE                             ***" << endl;
    cout << "****                                                ***" << endl;
    cout << "*******************************************************" << endl;
    ifstream ifs("C:\\WINDOWS\\system32\\kernel32.dll", ios_base::in|ios_base::binary);
    //ifstream ifs("D:\\3rd_lib\\zlib114dll\\dll32\\zlib.dll", ios_base::in|ios_base::binary);
    //ifstream ifs("D:\\3rd_lib\\libzplay-2.02-sdk\\libzplay.dll", ios_base::in|ios_base::binary);
    //ifstream ifs("kernel32.dll", ios_base::in|ios_base::binary);
    //D:\\Dropbox\\mini_opt\\codestat.exe");
    IMAGE_DOS_HEADER dos_header;
    dos_header.e_magic = 0;
    
    ifs.read(reinterpret_cast<char*>(&dos_header), sizeof(dos_header));
    char magic[3];
    magic[2] = 0;
    
    memcpy(magic, reinterpret_cast<char*>(&dos_header.e_magic), 2);
    if(strcmp(magic, "MZ") != 0){
        cout << "It's not a valid file without DOS header" << endl;
        return -1;
    }
    ifs.seekg(dos_header.e_lfanew, ios_base::beg);
    IMAGE_NT_HEADERS32 nt_header;
    ifs.read(reinterpret_cast<char*>(&nt_header), sizeof(nt_header));
    if ( nt_header.Signature  != 0x4550){
        cout << "It's not a valid file without PE magic" << endl;
    }
    
    //cout << nt_header.FileHeader;
    //cout << nt_header.OptionalHeader;

    vector<IMAGE_SECTION_HEADER> sec_headers;
    
    for(int i = 0; i < nt_header.FileHeader.NumberOfSections;  ++i){
        IMAGE_SECTION_HEADER sec_header;
        ifs.read(reinterpret_cast<char*>(&sec_header), sizeof(sec_header));
        cout << sec_header;
        sec_headers.push_back(sec_header);
    }

    //Export table
    const IMAGE_DATA_DIRECTORY* p_data_dir =
        &nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    
    if ( p_data_dir->VirtualAddress != 0 && p_data_dir->Size != 0){
        DWORD file_offset = 0;
        int delta_k = 0;
        
        for(vector<IMAGE_SECTION_HEADER>::const_iterator cit = sec_headers.begin(); cit != sec_headers.end(); ++cit){
            if (p_data_dir->VirtualAddress >= cit->VirtualAddress
                && p_data_dir->VirtualAddress < cit->VirtualAddress + cit->SizeOfRawData){

                delta_k =  cit->PointerToRawData -  cit->VirtualAddress;
                file_offset = p_data_dir->VirtualAddress + delta_k;
            }
        }
        if ( file_offset != 0){
            char* p_data = new char[p_data_dir->Size];
            ifs.seekg(file_offset, ios_base::beg);
            ifs.read(p_data, p_data_dir->Size);
            IMAGE_EXPORT_DIRECTORY* p_dir = reinterpret_cast<IMAGE_EXPORT_DIRECTORY*>(p_data);
            //print_ied(p_dir, file_offset, p_data_dir->VirtualAddress, p_data_dir->Size, delta_k);
            delete[] p_data;
        }
    }
    //Import table

    const IMAGE_DATA_DIRECTORY* p_data_dir2 =
        &nt_header.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    
    if ( p_data_dir2->VirtualAddress != 0 && p_data_dir2->Size != 0){
        DWORD file_offset = 0;
        int delta_k = 0;
        
        for(vector<IMAGE_SECTION_HEADER>::const_iterator cit = sec_headers.begin(); cit != sec_headers.end(); ++cit){
            if (p_data_dir2->VirtualAddress >= cit->VirtualAddress
                && p_data_dir2->VirtualAddress < cit->VirtualAddress + cit->SizeOfRawData){
                delta_k =  cit->PointerToRawData -  cit->VirtualAddress;
                cout << " Exist within : " << cit->Name << " deltaK  = " << delta_k << endl;        
                file_offset = p_data_dir2->VirtualAddress + delta_k;
                break;
            }
        }
        if ( file_offset != 0){
            char* p_data = new char[p_data_dir2->Size];
            ifs.seekg(file_offset, ios_base::beg);
            ifs.read(p_data, p_data_dir2->Size);
            IMAGE_IMPORT_DESCRIPTOR* p_iid = reinterpret_cast<IMAGE_IMPORT_DESCRIPTOR*>(p_data);
            int index = 0;
            while( p_iid[index].Characteristics != 0 || p_iid[index].TimeDateStamp != 0 ||
                   p_iid[index].ForwarderChain != 0 || p_iid[index].Name != 0 || p_iid[index].FirstThunk != 0 ){
                int name_file_offset = p_iid[index].Name + delta_k;
                
                cout << " Name = " << get_cstring(ifs, name_file_offset)
                     << " OriginalFirstThunk = " << hex << "0x" << p_iid[index].OriginalFirstThunk
                     << " Timedatastamp = " << hex << "0x" << p_iid[index].TimeDateStamp
                     << " ForwarderChain = " << hex << "0x" << p_iid[index].ForwarderChain
                     << " FirstThunk = " << hex << "0x" << p_iid[index].FirstThunk << endl;

                int original_first_thunk_offset = p_iid[index].OriginalFirstThunk + delta_k;
                vector<IMAGE_THUNK_DATA> thunk_datas = get_vector<IMAGE_THUNK_DATA>(ifs, original_first_thunk_offset);
                cout << "ThunkValue \thint \t ApiName" << endl;
                for(vector<IMAGE_THUNK_DATA>::const_iterator cit = thunk_datas.begin();
                    cit != thunk_datas.end(); ++cit){
                    WORD hint = get_data<WORD>(ifs, cit->u1.AddressOfData + delta_k);
                    string fun_name = get_cstring(ifs, cit->u1.AddressOfData + delta_k + sizeof(WORD));
                    
                    cout << hex << "0x" << cit->u1.AddressOfData
                         << hex << "\t0x" << hint
                         << "\t" << fun_name << endl;
                    
                }
                ++index;
            }
            delete[] p_data;
        }
    }
            
        
    return 0;
    
}
